import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import config

class StockDataManager:
    def __init__(self):
        """初始化股票数据管理器"""
        # AKShare不需要token
        
    def get_trade_date(self, offset=1):
        """获取交易日期，默认获取最近的一个交易日"""
        try:
            # 使用AKShare获取A股交易日历
            trade_cal = ak.tool_trade_date_hist_sina()
            # 转换为日期格式
            trade_cal['trade_date'] = pd.to_datetime(trade_cal['trade_date'])
            # 按日期降序排序
            trade_cal = trade_cal.sort_values('trade_date', ascending=False)
            
            # 获取交易日列表
            trade_dates = trade_cal['trade_date'].dt.strftime('%Y%m%d').tolist()
            
            # 返回指定偏移的交易日
            if offset < len(trade_dates):
                return trade_dates[offset]
            else:
                return trade_dates[-1]
        except Exception as e:
            # 如果出错，使用备用方法
            today = datetime.now()
            # 如果是周末，调整到周五
            if today.weekday() == 5:  # 星期六
                today = today - timedelta(days=1)
            elif today.weekday() == 6:  # 星期日
                today = today - timedelta(days=2)
                
            # 计算偏移日期
            target_date = today - timedelta(days=offset)
            # 如果是周末，调整到周五
            if target_date.weekday() == 5:  # 星期六
                target_date = target_date - timedelta(days=1)
            elif target_date.weekday() == 6:  # 星期日
                target_date = target_date - timedelta(days=2)
                
            return target_date.strftime('%Y%m%d')
    
    def get_stock_list(self):
        """获取A股股票列表"""
        try:
            # 使用AKShare获取A股股票列表
            stock_info = ak.stock_zh_a_spot_em()
            
            # 重命名列以保持一致性
            stock_info = stock_info.rename(columns={
                '代码': 'symbol',
                '名称': 'name',
                '行业': 'industry',
            })
            
            # 添加ts_code列（与原来的格式保持一致）
            stock_info['ts_code'] = stock_info['symbol'].apply(
                lambda x: x + ('.SH' if x.startswith('6') else '.SZ')
            )
            
            # 选择需要的列
            if 'industry' in stock_info.columns:
                result = stock_info[['ts_code', 'symbol', 'name', 'industry']]
            else:
                result = stock_info[['ts_code', 'symbol', 'name']]
                result['industry'] = '未知'
                
            return result
        except Exception as e:
            # 如果出错，返回空的DataFrame
            print(f"获取股票列表时出错: {str(e)}")
            return pd.DataFrame(columns=['ts_code', 'symbol', 'name', 'industry'])
    
    def get_daily_data(self, trade_date):
        """获取指定交易日的日线数据"""
        try:
            # 将日期格式转换为AKShare所需格式
            date_str = f"{trade_date[:4]}-{trade_date[4:6]}-{trade_date[6:8]}"
            
            # 使用AKShare获取当日所有A股数据
            df = ak.stock_zh_a_hist(symbol="000001", period="daily", 
                                   start_date=date_str, end_date=date_str, adjust="")
            
            # 如果当天数据不可用，尝试获取当日实时行情
            if df is None or df.empty:
                # 获取当日实时行情
                df = ak.stock_zh_a_spot_em()
            
            # 重命名列以保持一致性
            if df is not None and not df.empty:
                # 如果是历史数据
                if '日期' in df.columns and '开盘' in df.columns:
                    df = df.rename(columns={
                        '日期': 'trade_date',
                        '开盘': 'open',
                        '收盘': 'close',
                        '最高': 'high',
                        '最低': 'low',
                        '成交量': 'vol',
                        '成交额': 'amount',
                        '涨跌幅': 'pct_chg'
                    })
                # 如果是实时行情
                elif '代码' in df.columns and '最新价' in df.columns:
                    df = df.rename(columns={
                        '代码': 'symbol',
                        '名称': 'name',
                        '最新价': 'close',
                        '今开': 'open',  # 正确的开盘价列名
                        '最高': 'high',
                        '最低': 'low',
                        '成交量': 'vol',
                        '成交额': 'amount',
                        '涨跌幅': 'pct_chg',
                        '涨跌额': 'change',
                        '换手率': 'turnover_rate',
                        '昨收': 'pre_close'  # 添加昨收价列
                    })
                    # 添加交易日期列
                    df['trade_date'] = date_str.replace('-', '')
                
                # 添加ts_code列
                if 'symbol' in df.columns:
                    df['ts_code'] = df['symbol'].apply(
                        lambda x: f"{x}.{'SH' if x.startswith('6') else 'SZ'}"
                    )
                elif '代码' in df.columns:
                    df['ts_code'] = df['代码'].apply(
                        lambda x: f"{x}.{'SH' if x.startswith('6') else 'SZ'}"
                    )
                    df['symbol'] = df['代码']
                
                # 获取所有股票的数据
                if len(df) <= 1:  # 如果只有一支股票的数据
                    # 获取当日实时行情
                    all_stocks = ak.stock_zh_a_spot_em()
                    
                    # 重命名列
                    all_stocks = all_stocks.rename(columns={
                        '代码': 'symbol',
                        '名称': 'name',
                        '最新价': 'close',
                        '今开': 'open',  # 正确的开盘价列名
                        '最高': 'high',
                        '最低': 'low',
                        '成交量': 'vol',
                        '成交额': 'amount',
                        '涨跌幅': 'pct_chg',
                        '涨跌额': 'change',
                        '换手率': 'turnover_rate',
                        '昨收': 'pre_close'  # 添加昨收价列
                    })
                    
                    # 添加ts_code列和交易日期列
                    all_stocks['ts_code'] = all_stocks['symbol'].apply(
                        lambda x: f"{x}.{'SH' if x.startswith('6') else 'SZ'}"
                    )
                    all_stocks['trade_date'] = date_str.replace('-', '')
                    
                    df = all_stocks
            
            return df
        except Exception as e:
            print(f"获取日线数据时出错: {str(e)}")
            return pd.DataFrame()
    
    def get_top_volume_stocks(self, trade_date, top_n=100):
        """获取指定交易日成交量前N名的股票"""
        try:
            # 获取日线数据
            df = self.get_daily_data(trade_date)
            if df is None or df.empty:
                return pd.DataFrame()
                
            # 确保成交量列存在
            if 'vol' not in df.columns:
                if '成交量' in df.columns:
                    df['vol'] = df['成交量']
                elif 'volume' in df.columns:
                    df['vol'] = df['volume']
                else:
                    # 如果没有成交量列，返回空 DataFrame
                    print("无法获取成交量数据")
                    return pd.DataFrame()
            
            # 将成交量转换为数值类型
            df['vol'] = pd.to_numeric(df['vol'], errors='coerce')
            
            # 按成交量排序
            df = df.sort_values(by='vol', ascending=False)
            return df.head(top_n)
        except Exception as e:
            print(f"获取成交量前N名股票时出错: {str(e)}")
            return pd.DataFrame()
    
    def calculate_price_drop(self, prev_day_data, current_day_data):
        """计算股票价格跌幅
        
        Args:
            prev_day_data: 前一天的股票数据
            current_day_data: 当前的股票数据
            
        Returns:
            带有跌幅信息的DataFrame
        """
        try:
            # 检查数据是否为空
            if prev_day_data.empty or current_day_data.empty:
                return pd.DataFrame()
                
            # 确保两个DataFrame都有ts_code列
            if 'ts_code' not in prev_day_data.columns or 'ts_code' not in current_day_data.columns:
                raise ValueError("数据中缺少股票代码列")
            
            # 确保必要的列存在
            if 'close' not in prev_day_data.columns:
                if 'settlement' in prev_day_data.columns:
                    prev_day_data['close'] = prev_day_data['settlement']
                elif '收盘' in prev_day_data.columns:
                    prev_day_data['close'] = prev_day_data['收盘']
                elif '最新价' in prev_day_data.columns:
                    prev_day_data['close'] = prev_day_data['最新价']
                else:
                    raise ValueError("前一天数据中缺少收盘价列")
                
            if 'open' not in current_day_data.columns:
                if '开盘' in current_day_data.columns:
                    current_day_data['open'] = current_day_data['开盘']
                elif '开盘价' in current_day_data.columns:
                    current_day_data['open'] = current_day_data['开盘价']
                elif '今开' in current_day_data.columns:  # 添加对今开列的支持
                    current_day_data['open'] = current_day_data['今开']
                else:
                    # 如果没有开盘价，使用最新价作为替代
                    if '最新价' in current_day_data.columns:
                        print("警告: 当天数据中缺少开盘价列，使用最新价代替")
                        current_day_data['open'] = current_day_data['最新价']
                    else:
                        raise ValueError("当天数据中缺少开盘价列")
            
            # 创建数据的副本而不是视图，避免 SettingWithCopyWarning
            prev_day_data_copy = prev_day_data.copy()
            current_day_data_copy = current_day_data.copy()
            
            # 将数据转换为数值类型
            prev_day_data_copy['close'] = pd.to_numeric(prev_day_data_copy['close'], errors='coerce')
            current_day_data_copy['open'] = pd.to_numeric(current_day_data_copy['open'], errors='coerce')
            
            # 打印列名以进行调试
            print(f"prev_day_data 列名: {prev_day_data_copy.columns.tolist()}")
            print(f"current_day_data 列名: {current_day_data_copy.columns.tolist()}")
            
            # 合并前一天和当天的数据
            merged_data = pd.merge(
                prev_day_data_copy[['ts_code', 'close']], 
                current_day_data_copy[['ts_code', 'open']], 
                on='ts_code', 
                suffixes=('_prev', '_current')
            )
            
            # 打印合并后的列名以进行调试
            print(f"merged_data 列名: {merged_data.columns.tolist()}")
            
            # 计算跌幅百分比: (昨收 - 今开) / 昨收 * 100
            # 确保使用正确的列名
            merged_data['drop_percent'] = (merged_data['close_prev'] - merged_data['open']) / merged_data['close_prev'] * 100
            
            # 获取完整的当天数据
            result = pd.merge(current_day_data_copy, merged_data[['ts_code', 'drop_percent']], on='ts_code')
            
            return result
        except Exception as e:
            print(f"计算跌幅时出错: {str(e)}")
            return pd.DataFrame()
    
    def get_top_price_drop_stocks(self, top_volume=100, top_drop=10):
        """获取昨日成交量前N名中今日开盘价跌幅最大的M只股票"""
        try:
            # 获取最近两个交易日
            yesterday = self.get_trade_date(offset=1)
            today = self.get_trade_date(offset=0)
            
            if not yesterday or not today:
                return pd.DataFrame(), "无法获取交易日数据"
            
            print(f"正在获取交易日数据: 昨日={yesterday}, 今日={today}")
            
            # 获取昨日成交量前N名的股票
            top_volume_stocks = self.get_top_volume_stocks(yesterday, top_volume)
            if top_volume_stocks.empty:
                return pd.DataFrame(), "无法获取昨日成交量数据"
                
            print(f"成功获取昨日成交量前{top_volume}名的股票数据")
                
            # 获取今日的股票数据
            today_data = self.get_daily_data(today)
            if today_data.empty:
                return pd.DataFrame(), "无法获取今日股票数据"
                
            print(f"成功获取今日股票数据")
                
            # 只保留昨日成交量前N名的今日数据
            today_top_volume_data = today_data[today_data['ts_code'].isin(top_volume_stocks['ts_code'])]
            if today_top_volume_data.empty:
                return pd.DataFrame(), "无法匹配昨日成交量前N名的今日数据"
                
            print(f"成功匹配昨日成交量前{top_volume}名的今日数据")
                
            # 计算跌幅
            stocks_with_drop = self.calculate_price_drop(top_volume_stocks, today_top_volume_data)
            if stocks_with_drop.empty:
                return pd.DataFrame(), "计算跌幅时出错"
                
            print("成功计算股票跌幅")
                
            # 按跌幅排序（降序，因为跌幅是正值）
            sorted_stocks = stocks_with_drop.sort_values(by='drop_percent', ascending=False)
            
            # 获取跌幅最大的M只股票
            top_drop_stocks = sorted_stocks.head(top_drop)
            
            print(f"成功获取跌幅前{top_drop}名的股票")
            
            # 获取股票的基本信息
            stock_info = self.get_stock_list()
            
            # 确俚stock_info不为空
            if stock_info.empty:
                # 如果无法获取股票基本信息，直接返回跌幅数据
                print("无法获取股票基本信息，直接返回跌幅数据")
                return top_drop_stocks, None
                
            print("成功获取股票基本信息")
                
            # 合并股票信息
            # 确保必要的列存在
            if 'name' not in stock_info.columns:
                if '名称' in stock_info.columns:
                    stock_info['name'] = stock_info['名称']
                else:
                    stock_info['name'] = '未知'
                
            if 'industry' not in stock_info.columns:
                if '行业' in stock_info.columns:
                    stock_info['industry'] = stock_info['行业']
                else:
                    # 如果没有行业信息，添加一个空列
                    stock_info['industry'] = '未知'
                
            # 合并数据
            result = pd.merge(top_drop_stocks, stock_info[['ts_code', 'name', 'industry']], on='ts_code', how='left')
            
            print("成功合并股票信息和跌幅数据")
            return result, None
            
        except Exception as e:
            error_msg = f"获取数据时出错: {str(e)}"
            print(error_msg)
            return pd.DataFrame(), error_msg
    
    def format_stock_data(self, data):
        """格式化股票数据用于显示"""
        if data.empty:
            return pd.DataFrame()
            
        # 选择要显示的列
        display_columns = [
            'ts_code', 'name', 'industry', 'open', 'close', 'high', 'low', 
            'vol', 'amount', 'drop_percent'
        ]
        
        # 确保所有列都存在
        existing_columns = [col for col in display_columns if col in data.columns]
        formatted_data = data[existing_columns].copy()
        
        # 格式化列名
        column_names = {
            'ts_code': '股票代码',
            'name': '股票名称',
            'industry': '所属行业',
            'open': '开盘价',
            'close': '收盘价',
            'high': '最高价',
            'low': '最低价',
            'vol': '成交量',
            'amount': '成交额',
            'drop_percent': '跌幅(%)'
        }
        
        # 重命名列
        formatted_data = formatted_data.rename(columns=column_names)
        
        # 格式化数值
        if '跌幅(%)' in formatted_data.columns:
            formatted_data['跌幅(%)'] = formatted_data['跌幅(%)'].round(2)
        
        return formatted_data
